from langchain.agents import create_tool_calling_agent, AgentExecutor, create_react_agent
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.tools import tool

from src.ai.langchain.init_llm import get_llm

llm = get_llm()


@tool("get_weather")
def get_weather(city: str) -> str:
    """
    获取指定城市的天气情况。

    参数：
    city (str): 城市名称

    返回：
    str: 当前天气状况
    """
    return f"{city}今天有雨"


@tool("get_temperature")
def get_temperature(city: str) -> str:
    """
    获取指定城市的温度情况。

    参数：
    city (str): 城市名称

    返回：
    str: 当前温度
    """
    return f"{city}当前的气温是25度"


@tool("get_hobby")
def get_hobby(name: str) -> str:
    """
    查询用户的兴趣爱好。

    参数：
    name (str): 用户名

    返回：
    str: 该用户的兴趣爱好
    """
    return f"{name}的爱好是看电影"


message = """
    你是一个智能体，擅长使用工具解决问题。
    你可以使用以下工具：
    {tools}
    
    请使用如下格式：

    Question: {input}
    {agent_scratchpad}
    
    Thought: 你在思考什么？
    Action: 工具名称（{tool_names} 中的一个）
    Action Input: 工具的输入
    Observation: 工具返回的结果
    ...（可以重复 Thought/Action/Action Input/Observation 多次）
    Thought: 我知道最终答案了
    Final Answer: 最终答案是什么
    
    现在开始！
    
    Question: {input}
    {agent_scratchpad}
"""

prompt = PromptTemplate.from_template(message)

tools = [
    get_hobby,
    get_weather,
    get_temperature,
]

agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)

# agent负责“思考”应该调用哪个工具，agent_executor负责执行tool，所以执行器里面也需要传入工具列表
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

print(agent_executor.invoke({"input": "天津的温度怎么样"}))
